/*
 * Description: performance common utilities.
 */
#include "perf_utils.h"
#include <algorithm>
#include <sstream>
#include "Windows.h"
#include <TlHelp32.h>
#include <sddl.h>
#include "log.h"
#include "performance/perf_constants.h"
#include "text_util/string_proc.h"
#include "common/guest_utils.h"

namespace PERF {

constexpr const DWORD TIME_WAIT_FOR_PROCESS_EXIT = 200;

std::vector<std::string> PerfUtils::Split(const std::string &str, const std::string &delim)
{
    std::size_t previous = 0;
    std::size_t current = str.find_first_of(delim);
    std::vector<std::string> elems;
    while (current != std::string::npos) {
        if (current > previous) {
            elems.push_back(str.substr(previous, current - previous));
        }
        previous = current + delim.size();
        current = str.find_first_of(delim, previous);
    }
    if (previous != str.size()) {
        elems.push_back(str.substr(previous));
    }
    return elems;
}

std::vector<std::wstring> PerfUtils::SplitUnicode(const std::wstring &str, const std::wstring &delim)
{
    std::wstring s = str;
    std::vector<std::wstring> res;
    while (!s.empty()) {
        size_t pos = s.find(delim);
        if (pos == std::wstring::npos) {
            res.emplace_back(s);
            break;
        }
        res.emplace_back(s.substr(0, pos));
        s = s.substr(pos + delim.size());
    }
    return res;
}

void PerfUtils::ToLower(std::vector<std::string> &strs)
{
    for (std::string &str : strs) {
        std::transform(str.begin(), str.end(), str.begin(), tolower);
    }
}

bool PerfUtils::GetUserToken(HANDLE *phToken)
{
    if (phToken == nullptr) {
        return false;
    }
    DWORD pidExplorer = GetProcessIdByName(PNAME_EXPLORER.c_str());
    HANDLE hExplorer = OpenProcess(PROCESS_ALL_ACCESS, false, pidExplorer);
    if (!hExplorer) {
        return false;
    }
    if (!OpenProcessToken(hExplorer, TOKEN_ALL_ACCESS, phToken)) {
        CloseHandle(hExplorer);
        return false;
    }
    CloseHandle(hExplorer);
    return true;
}

bool PerfUtils::LookupCurrentUserSid(HANDLE &hToken, SID *&sid)
{
    DWORD size = 0;
    TOKEN_USER *user;

    GetTokenInformation(hToken, TokenUser, NULL, 0, &size);
    if (!size) {
        return false;
    }

    user = (TOKEN_USER *)malloc(size);
    if (user == nullptr) {
        return false;
    }
    if (!GetTokenInformation(hToken, TokenUser, user, size, &size)) {
        free(user);
        user = nullptr;
        return false;
    }
    size = GetLengthSid(user->User.Sid);
    sid = (SID *)malloc(size);
    if (sid == nullptr) {
        free(user);
        user = nullptr;
        return false;
    }

    bool ret = CopySid(size, sid, user->User.Sid);
    free(user);
    user = nullptr;
    return ret;
}

std::string PerfUtils::GetCurrentUserSid()
{
    HANDLE hToken = INVALID_HANDLE_VALUE;
    if (!GetUserToken(&hToken)) {
        return userSid;
    }

    SID *sid = nullptr;
    LookupCurrentUserSid(hToken, sid);
    CloseHandle(hToken);
    hToken = INVALID_HANDLE_VALUE;
    if (sid == nullptr) {
        return userSid;
    }

    LPSTR lpSidCStr = nullptr;
    int ret = ConvertSidToStringSid(sid, &lpSidCStr);
    free(sid);
    sid = nullptr;
    if (lpSidCStr == nullptr) {
        return userSid;
    }

    std::string sidStr = lpSidCStr;
    LocalFree(lpSidCStr);
    lpSidCStr = nullptr;
    if (!sidStr.empty() && userSid != sidStr) {
        userSid = sidStr;
    }
    return userSid;
}

DWORD PerfUtils::GetProcessIdByName(LPCTSTR pszExeFile)
{
    DWORD nPid = 0;
    PROCESSENTRY32 pe = { sizeof PROCESSENTRY32 };
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);
    if (hSnapshot != INVALID_HANDLE_VALUE && Process32First(hSnapshot, &pe)) {
        while (Process32Next(hSnapshot, &pe)) {
            std::string argExeFile = pszExeFile;
            std::string itemExeFile = pe.szExeFile;
            std::transform(argExeFile.begin(), argExeFile.end(), argExeFile.begin(), ::toupper);
            std::transform(itemExeFile.begin(), itemExeFile.end(), itemExeFile.begin(), ::toupper);
            if (argExeFile == itemExeFile) {
                nPid = pe.th32ProcessID;
                break;
            }
        }
    }
    if (hSnapshot != INVALID_HANDLE_VALUE) {
        CloseHandle(hSnapshot);
        hSnapshot = INVALID_HANDLE_VALUE;
    }
    return nPid;
}

bool PerfUtils::QueryProcessNameByPid(IN DWORD pid, OUT std::wstring &processFullName)
{
    if ((pid == SYSTEM_IDLE_PROCESS) || (pid == SYSTEM_SYSTEM_PROCESS)) {
        return false;
    }
    HANDLE procHandle = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);
    if (procHandle == nullptr) {
        procHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid);
        if (procHandle == nullptr) {
            return false;
        }
    }
    WCHAR processName[MAX_PATH] = { 0 };
    DWORD size = sizeof(processName) / sizeof(WCHAR) - 1;
    BOOL ret = QueryFullProcessImageNameW(procHandle, 0, processName, &size);
    if (ret == false) {
        CloseHandle(procHandle);
        return false;
    } else {
        processFullName = processName;
    }
    CloseHandle(procHandle);
    return true;
}

bool PerfUtils::QueryProcessNameByPid(IN DWORD pid, OUT std::string &processFullName)
{
    if ((pid == SYSTEM_IDLE_PROCESS) || (pid == SYSTEM_SYSTEM_PROCESS)) {
        return false;
    }
    HANDLE procHandle = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);
    if (procHandle == nullptr) {
        procHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid);
        if (procHandle == nullptr) {
            return false;
        }
    }
    char processName[MAX_PATH] = { 0 };
    DWORD size = sizeof(processName) / sizeof(char) - 1;
    BOOL ret = QueryFullProcessImageNameA(procHandle, 0, processName, &size);
    if (ret == false) {
        CloseHandle(procHandle);
        return false;
    } else {
        processFullName = processName;
    }
    CloseHandle(procHandle);
    return true;
}

bool PerfUtils::IsProcessExists(IN DWORD processId)
{
    HANDLE process = OpenProcess(PROCESS_QUERY_INFORMATION, false, processId);
    if (process == nullptr) {
        process = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, processId);
        if (process == nullptr) {
            return false;
        }
    }
    CloseHandle(process);
    return true;
}

bool PerfUtils::IsProcessExists(IN const std::wstring &processName)
{
    return GetProcessIdByName((StringProc::WStringToString(processName)).c_str()) != 0;
}

bool PerfUtils::IsProcessExists(IN const std::string &processName)
{
    return GetProcessIdByName(processName.c_str()) != 0;
}

bool PerfUtils::GetFilepathFromCmdLine(const std::string &cmdLine, std::string &filepath)
{
    std::vector<std::string> cmdLineParts = PERF_UTILS.Split(cmdLine, "\"");
    int idx = 0;
    for (std::string &cmdLinePart : cmdLineParts) {
        if (idx++ == 0) {
            continue;
        }
        std::string cmdLinePartLower = cmdLinePart;
        std::transform(cmdLinePartLower.begin(), cmdLinePartLower.end(), cmdLinePartLower.begin(), ::tolower);
        if (cmdLinePartLower.find("\\") == std::string::npos) {
            continue;
        }
        filepath = cmdLinePart;
        return true;
    }
    return false;
}
}  // namespace PERF